
from monai.networks.nets.swin_unetr import SwinUNETR

import torch.nn as nn
from medical.loss.dice_ce_loss import EntropyAndDiceLoss

class SwinUNETRWrapper(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 img_size,
                 feature_size=48,
                 depths=(2,2,2,2),
                 num_heads=(3,6,12,24),
                 drop_rate=0.0,
                 attn_drop_rate=0.0):
        super().__init__()
        self.model = SwinUNETR(spatial_dims=3,
                               img_size=img_size,
                               in_channels=in_channels,
                               out_channels=out_channels,
                               depths=depths,
                               feature_size=feature_size,
                               num_heads=num_heads,
                               drop_rate=drop_rate,
                               attn_drop_rate=attn_drop_rate)

        self.dice_ce_loss = EntropyAndDiceLoss()

    def forward(self, x):
        if type(x) is dict:
            return self.forward_train(x)

        return self.model(x)

    def forward_train(self, data):
        x = data["image"]
        label = data.get("label", None)

        logits = self.forward(x)

        return_data = {"logits": logits}
        if label is not None:
            loss = self.compute_loss(logits, label)
            return_data["loss"] = loss

        return return_data

    def compute_loss(self, pred, label):
        return self.dice_ce_loss(pred, label)